package com.sdses.flink.udf;

import org.apache.flink.table.functions.FunctionContext;
import org.apache.flink.table.functions.ScalarFunction;
import org.geotools.data.DataStore;
import org.geotools.data.DataStoreFinder;
import org.geotools.data.simple.SimpleFeatureIterator;
import org.geotools.data.simple.SimpleFeatureSource;
import org.geotools.geojson.geom.GeometryJSON;
import org.geotools.geometry.DirectPosition2D;
import org.geotools.geometry.jts.WKBReader;
import org.geotools.jdbc.JDBCDataStoreFactory;
import org.geotools.referencing.GeodeticCalculator;
import org.locationtech.jts.geom.Coordinate;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.io.ParseException;
import org.locationtech.jts.io.WKBHexFileReader;
import org.locationtech.jts.io.WKTReader;
import org.locationtech.jts.operation.distance.DistanceOp;
import org.opengis.feature.simple.SimpleFeature;
import org.opengis.referencing.operation.TransformException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.Reader;
import java.io.StringReader;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author jiwei
 * @description
 * @date 2023/6/5 14:51
 */
public class UDFGeomDistance extends ScalarFunction {
    private final static Logger log = LoggerFactory.getLogger(UDFGeomDistance.class);
    private static final String TYPE_WKT = "WKT";
    private static final String TYPE_GEO_JSON = "GeoJSON";

    private List<Map<String,String>> geoList = new ArrayList<>();

    /**
     * 计算两个点的距离
     * @param sourceString
     * @return distance
     */
    public List<Long> eval(String sourceString, Integer distance) {
        log.info("UDFGeoLookupDistance param:{},{},{}",sourceString);

        Geometry sourceGeometry = toGeometry(sourceString, TYPE_WKT);
        List<Long> list = geoList.stream().filter(entity -> {
            Geometry targetGeometry = toGeometry(entity.get("point_wkt"), TYPE_WKT);

            // 首先找出最近的两个点
            Coordinate[] points = DistanceOp.nearestPoints(sourceGeometry,targetGeometry);
            Coordinate point0 = points[0];
            Coordinate point1 = points[1];

            // 创建大地计算器
            GeodeticCalculator calculator = new GeodeticCalculator();

            // 计算距离
            try {
                calculator.setStartingPosition(new DirectPosition2D(point0.x, point0.y));
                calculator.setDestinationPosition(new DirectPosition2D(point1.x, point1.y));
            } catch (TransformException e) {
                log.error("距离计算失败：",e);
                throw new RuntimeException(e.toString());
            }
            double distanceResult = calculator.getOrthodromicDistance();

            return distanceResult <= distance;
        }).map(entity -> Long.parseLong(entity.get("poi_id"))).collect(Collectors.toList());

        return list;
    }

    @Override
    public void open(FunctionContext context) throws Exception {
        DataStore dataStore = getDataStore();
        geoList = getValueList(dataStore);
    }

    private Geometry toGeometry(String geoString, String type){
        Geometry geometry = null;
        try {
            if(TYPE_WKT.equals(type)) {
                WKTReader reader = new WKTReader();
                geometry = reader.read(geoString);
            }else if (TYPE_GEO_JSON.equals(type)){
                GeometryJSON geometryJSON = new GeometryJSON();
                Reader reader = new StringReader(geoString);
                geometry = geometryJSON.read(reader);
            }
        } catch (Exception e) {
            log.error("WKT、GeoJSON转换失败:",e);
            throw new RuntimeException("WKT、GeoJSON转换失败!");
        }

        return geometry;
    }

    /**
     * 初始化数据源
     * @param dataStore
     * @return
     */
    private static List<Map<String,String>> getValueList(DataStore dataStore){
        SimpleFeatureIterator iterator = null;
        List<Map<String,String>> geoList = new ArrayList<>();
        try {
            SimpleFeatureSource fSource = dataStore.getFeatureSource("poi_info_wkt");
            iterator = fSource.getFeatures().features();
        } catch (Exception e) {
            log.error("获取表失败[poi_info_wkt]:", e);
            throw new RuntimeException(e.toString());
        }
        while (iterator.hasNext()) {

            // 获取每一个要素
            SimpleFeature feature = iterator.next();
            Map<String,String> map = new HashMap<>();
            map.put("poi_id", String.valueOf(feature.getAttribute("poi")));
            map.put("point_wkt", String.valueOf(feature.getAttribute("point_wkt_")));
            geoList.add(map);
        }

        return geoList;
    }

    private static DataStore getDataStore(){
        Map<String, Object> params = new HashMap<>();

        // 需要连接何种数据库，postgis or mysql
        params.put(JDBCDataStoreFactory.DBTYPE.key, "mysql");
        params.put(JDBCDataStoreFactory.HOST.key, "192.168.102.155");
        params.put(JDBCDataStoreFactory.PORT.key, new Integer(33066));
        params.put(JDBCDataStoreFactory.DATABASE.key, "db_datacube_test");
        params.put(JDBCDataStoreFactory.USER.key, "datacube_test");
        params.put(JDBCDataStoreFactory.PASSWD.key, "Gbd#XERA@19d1ZtD");
        DataStore dataStore;
        try {
            dataStore = DataStoreFinder.getDataStore(params);
        } catch (Exception e) {
            log.error("数据库连接失败：", e);
            throw new RuntimeException(e.toString());
        }
        return dataStore;
    }

    public static void main(String[] args) throws ParseException, IOException {

        String geoString = "{ \"type\": \"MultiLineString\", \"coordinates\": [ [ [ 117.109982198520697, 36.728969675411065 ], [ 117.109901936726033, 36.729570852413815 ] ] ] }";
        String type = "json";
        String wkbString = "{\"wkb\":\"AQEAAAAAAAAAAAAkQAAAAAAAADRA\",\"srid\":null}";
//        String wkbString = "AQEAAAAAAAAAAAAkQAAAAAAAADRA";

//        WKBReader reader = new WKBReader();
//        new WKBHexFileReader()
//        Geometry geometry = reader.read(wkbString.getBytes(StandardCharsets.UTF_8));
        GeometryJSON geometryJson = new GeometryJSON();
        Geometry geometry = geometryJson.read(geoString);

        WKBReader wkbReader = new WKBReader();
        System.out.println(geometry);


        byte[] bytes = org.locationtech.jts.io.WKBReader.hexToBytes(wkbString);
        WKBReader reader = new WKBReader();
        Geometry geometry1 = reader.read(bytes);
        System.out.println(geometry1);


    }

}
